#include <iostream>

namespace shr
{
    enum color
    {
        RED,
        BLACK
    };

    template <class T>
    struct rbTreeNode
    {
        T _data;
        color _color;
        rbTreeNode *_parent;
        rbTreeNode *_leftChild;
        rbTreeNode *_rightChild;

        rbTreeNode(T data = T())
            : _data(data), _color(RED),
              _parent(nullptr),
              _leftChild(nullptr),
              _rightChild(nullptr)
        { }
    };

    template <class T>
    class rb_tree
    {
    private:
        typedef rbTreeNode<T> Node;

    public:
        rb_tree()
            : _root(nullptr)
        { }

        bool insert(const T &data)
        {
            Node *newNode = new Node(data);
            if (_root == nullptr)
            {
                _root = newNode;
                _root->_color = BLACK; //根的颜色为黑
                return true;
            }
            Node *curNode = _root;
            Node *parent = nullptr;
            while (curNode)
            {
                if (data < curNode->_data)
                {
                    parent = curNode;
                    curNode = curNode->_leftChild;
                }
                else if (data > curNode->_data)
                {
                    parent = curNode;
                    curNode = curNode->_rightChild;
                }
                else
                {
                    return false;
                }
            }
            if (data < parent->_data)
            {
                parent->_leftChild = newNode;
                newNode->_parent = parent;
            }
            else // if(data > parent->_data)
            {
                parent->_rightChild = newNode;
                newNode->_parent = parent;
            }
            curNode = newNode;

            while (parent && parent->_color == RED)
            {
                Node *ppNode = parent->_parent;
                if (parent == ppNode->_leftChild)
                {
                    Node *uncleNode = ppNode->_rightChild;
                    if (uncleNode && uncleNode->_color == RED)
                    {
                        parent->_color = uncleNode->_color = BLACK;
                        ppNode->_color = RED;
                        curNode = ppNode;
                        parent = curNode->_parent;
                    }
                    else
                    {
                        if (curNode == parent->_leftChild)
                        {
                            rotateRight(ppNode);
                            parent->_color = BLACK;
                            curNode->_color = ppNode->_color = RED;
                        }
                        else if (curNode == parent->_rightChild)
                        {
                            rotateLeftRight(ppNode);
                            curNode->_color = BLACK;
                            parent->_color = ppNode->_color = RED;
                        }
                        else {
                            std::cout << __FILE__ << ":" << __LINE__ << ":insert err" << std::endl;
                        }
                        break;
                    }
                }
                else if (parent == ppNode->_rightChild)
                {
                    Node *uncleNode = ppNode->_leftChild;
                    if (uncleNode && uncleNode->_color == RED)
                    {
                        parent->_color = uncleNode->_color = BLACK;
                        ppNode->_color = RED;
                        curNode = ppNode;
                        parent = curNode->_parent;
                    }
                    else
                    {
                        if(curNode == parent->_rightChild)
                        {
                            rotateLeft(ppNode);
                            parent->_color = BLACK;
                            uncleNode->_color = parent->_color = RED;
                        }
                        else if(curNode == parent->_leftChild)
                        {
                            rotateRightLeft(ppNode);
                            curNode->_color = BLACK;
                            ppNode->_color = parent->_color = RED;
                        }
                        else {
                            std::cout << __FILE__ << ":" << __LINE__ << ":insert err" << std::endl;
                        }
                        break;
                    }
                }
                else {
                    std::cout << __FILE__ << ":" << __LINE__ << ":insert err" << std::endl;
                }
            }
            _root->_color = BLACK;
            return true;
        }

        void printTest()
        {
            _preorder_travel(_root);
        }

        ~rb_tree()
        {
            _destroy(_root);
        }

    private:
        void _preorder_travel(Node* root)
        {
            if(root == nullptr) { return; }
            std::cout << root->_data << ' ';
            _preorder_travel(root->_leftChild);
            _preorder_travel(root->_rightChild);
        }

        void rotateLeft(Node *parent)
        {
            // std::cout << "rotateLeft, val:" << parent->_data  << std::endl;
            Node *curNode = parent->_rightChild;
            Node *curLeft = curNode->_leftChild;

            Node *ppnode = parent->_parent;

            parent->_rightChild = curLeft;
            if (curLeft != nullptr)
            {
                curLeft->_parent = parent;
            }
            curNode->_leftChild = parent;
            parent->_parent = curNode;

            if (ppnode == nullptr)
            {
                _root = curNode;
            }
            else
            {
                if (parent == ppnode->_leftChild)
                {
                    ppnode->_leftChild = curNode;
                }
                else if (parent == ppnode->_rightChild)
                {
                    ppnode->_rightChild = curNode;
                }
                else
                {
                    perror("insert err");
                    exit(1);
                }
            }
            curNode->_parent = ppnode;
        }

        void rotateRight(Node *parent)
        {
            // std::cout << "rotateRight" << std::endl;
            Node *curNode = parent->_leftChild;
            Node *curRight = curNode->_rightChild;
            Node *ppNode = parent->_parent;

            parent->_leftChild = curRight;
            if (curRight != nullptr)
            {
                curRight->_parent = parent;
            }
            curNode->_rightChild = parent;
            parent->_parent = curNode;

            if (ppNode == nullptr)
            {
                _root = curNode;
            }
            else
            {
                if (parent == ppNode->_leftChild)
                {
                    ppNode->_leftChild = curNode;
                }
                else if (parent == ppNode->_rightChild)
                {
                    ppNode->_rightChild = curNode;
                }
                else
                {
                    perror("insert err");
                    exit(1);
                }
            }
            curNode->_parent = ppNode;
        }

        void rotateLeftRight(Node *parent)
        {
            // std::cout << "rotateLeftRight, val:" << parent->_data << std::endl;
            Node *curNode = parent->_leftChild;
            Node *curRight = curNode->_rightChild;

            rotateLeft(curNode);
            rotateRight(parent);
        }

        void rotateRightLeft(Node *parent)
        {
            // std::cout << "rotateRightLeft, val:" << parent->_data << std::endl;
        }

        void _destroy(Node *root)
        {
            if (root == nullptr)
            {
                return;
            }
            _destroy(root->_leftChild);
            _destroy(root->_rightChild);
            delete root;
        }

    private:
        Node *_root;
    };
}